{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "fdc17041-bd7a-4c73-94a3-98ea6ef8407b",
   "metadata": {},
   "outputs": [],
   "source": [
    "#卷积神经网络、手写字母识别"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "1657039b-5432-4818-bec2-88bbeb013e9f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test set size: 10000\n",
      "Training set size: 60000\n",
      "Number of training samples: 43360\n",
      "Number of cross-validation samples: 10850\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "label_size = 18 # Label size\n",
    "ticklabel_size = 14 # Tick label size\n",
    "    \n",
    "# Define a transform to normalize the data\n",
    "transform = transforms.Compose([\n",
    "    transforms.ToTensor()\n",
    "])\n",
    "\n",
    "# Load test data from the MNIST\n",
    "testset = torchvision.datasets.MNIST(root='./Data', train=False, download=False, transform=transform)\n",
    "print(f\"Test set size: {len(testset)}\")\n",
    "\n",
    "# Load training data from the MNIST\n",
    "trainset = torchvision.datasets.MNIST(root='./Data', train=True, download=False, transform=transform)\n",
    "print(f\"Training set size: {len(trainset)}\")\n",
    "\n",
    "# Rate of trX and cvX\n",
    "tr_cv_rate = 0.8\n",
    "\n",
    "# Create a list to store indices for each class unique()\n",
    "class_indices = [[] for _ in range(10)]  # 10 classes in MNIST\n",
    "\n",
    "# Populate class_indices\n",
    "for idx, (_, label) in enumerate(trainset):\n",
    "    class_indices[label].append(idx)\n",
    "\n",
    "# Calculate the number of samples for each class in training and validation sets\n",
    "train_size_per_class = int(tr_cv_rate * min(len(indices) for indices in class_indices))\n",
    "val_size_per_class = min(len(indices) for indices in class_indices) - train_size_per_class\n",
    "\n",
    "# Create balanced train and validation sets\n",
    "train_indices = []\n",
    "val_indices = []\n",
    "for indices in class_indices:\n",
    "    train_indices.extend(indices[:train_size_per_class])\n",
    "    val_indices.extend(indices[train_size_per_class:train_size_per_class + val_size_per_class])\n",
    "\n",
    "# Create Subset datasets\n",
    "from torch.utils.data import Subset\n",
    "trX = Subset(trainset, train_indices)\n",
    "cvX = Subset(trainset, val_indices)\n",
    "\n",
    "print(f\"Number of training samples: {len(trX)}\")\n",
    "print(f\"Number of cross-validation samples: {len(cvX)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "085e3a50-f0da-4ff5-8f92-ddeabdf1fc27",
   "metadata": {},
   "outputs": [],
   "source": [
    "#构建DataLoaders，准备训练模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "a92b9b3e-e44d-4dfd-b573-5d5d1993f813",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "image_channels is 1\n",
      "tensor([[0., 0., 0., 1., 0., 0., 0., 0., 0., 0.],\n",
      "        [0., 0., 0., 0., 0., 0., 0., 1., 0., 0.],\n",
      "        [0., 0., 1., 0., 0., 0., 0., 0., 0., 0.],\n",
      "        [0., 0., 0., 0., 0., 1., 0., 0., 0., 0.],\n",
      "        [0., 0., 0., 0., 0., 1., 0., 0., 0., 0.],\n",
      "        [1., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
      "        [0., 0., 1., 0., 0., 0., 0., 0., 0., 0.],\n",
      "        [0., 0., 0., 0., 0., 0., 0., 0., 1., 0.],\n",
      "        [0., 0., 0., 0., 0., 0., 0., 1., 0., 0.],\n",
      "        [0., 0., 0., 0., 0., 0., 0., 0., 1., 0.],\n",
      "        [0., 0., 0., 0., 0., 0., 1., 0., 0., 0.],\n",
      "        [0., 0., 0., 0., 0., 0., 0., 1., 0., 0.],\n",
      "        [0., 0., 0., 0., 0., 0., 0., 0., 0., 1.],\n",
      "        [0., 0., 0., 0., 0., 0., 0., 0., 1., 0.],\n",
      "        [0., 0., 0., 0., 0., 1., 0., 0., 0., 0.],\n",
      "        [0., 0., 0., 0., 1., 0., 0., 0., 0., 0.],\n",
      "        [0., 0., 0., 0., 0., 0., 0., 0., 0., 1.],\n",
      "        [0., 0., 0., 0., 0., 0., 0., 0., 0., 1.],\n",
      "        [0., 0., 0., 0., 0., 0., 0., 0., 0., 1.],\n",
      "        [0., 0., 1., 0., 0., 0., 0., 0., 0., 0.],\n",
      "        [0., 0., 0., 0., 1., 0., 0., 0., 0., 0.],\n",
      "        [0., 0., 1., 0., 0., 0., 0., 0., 0., 0.],\n",
      "        [1., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
      "        [0., 0., 0., 1., 0., 0., 0., 0., 0., 0.],\n",
      "        [0., 0., 0., 0., 0., 0., 0., 0., 0., 1.],\n",
      "        [0., 0., 0., 0., 0., 0., 0., 0., 1., 0.],\n",
      "        [0., 0., 0., 1., 0., 0., 0., 0., 0., 0.],\n",
      "        [0., 0., 0., 0., 0., 0., 1., 0., 0., 0.],\n",
      "        [0., 0., 1., 0., 0., 0., 0., 0., 0., 0.],\n",
      "        [0., 0., 0., 1., 0., 0., 0., 0., 0., 0.],\n",
      "        [0., 0., 0., 0., 0., 0., 1., 0., 0., 0.],\n",
      "        [1., 0., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
      "        [0., 0., 0., 0., 1., 0., 0., 0., 0., 0.],\n",
      "        [0., 1., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
      "        [0., 0., 1., 0., 0., 0., 0., 0., 0., 0.],\n",
      "        [0., 0., 0., 0., 0., 0., 0., 1., 0., 0.],\n",
      "        [0., 0., 1., 0., 0., 0., 0., 0., 0., 0.],\n",
      "        [0., 0., 0., 0., 0., 1., 0., 0., 0., 0.],\n",
      "        [0., 1., 0., 0., 0., 0., 0., 0., 0., 0.],\n",
      "        [0., 0., 0., 0., 0., 0., 0., 0., 0., 1.],\n",
      "        [0., 0., 0., 0., 0., 0., 0., 1., 0., 0.],\n",
      "        [0., 0., 0., 0., 0., 1., 0., 0., 0., 0.]])\n"
     ]
    }
   ],
   "source": [
    "batch_size = 42 # Define training batch 1，\n",
    "\n",
    "def one_hot_collate(batch):\n",
    "    data = torch.stack([item[0] for item in batch])\n",
    "    labels = torch.tensor([item[1] for item in batch])\n",
    "    one_hot_labels = torch.zeros(labels.size(0), 10)  # 10 classes in MNIST 【0，1，0，0】\n",
    "    one_hot_labels.scatter_(1, labels.unsqueeze(1), 1)\n",
    "    return data, one_hot_labels\n",
    "\n",
    "trLoader = torch.utils.data.DataLoader(trX, batch_size=batch_size, shuffle=True, num_workers=0, collate_fn=one_hot_collate)\n",
    "cvLoader = torch.utils.data.DataLoader(cvX, batch_size=batch_size, shuffle=False, num_workers=0, collate_fn=one_hot_collate)\n",
    "teLoader = torch.utils.data.DataLoader(testset, batch_size=batch_size, shuffle=False, num_workers=0, collate_fn=one_hot_collate)\n",
    "\n",
    "# Get a batch of training data\n",
    "dataiter = iter(trLoader)\n",
    "data, labels = next(dataiter)\n",
    "\n",
    "image_channels = data[0].numpy().shape[0]\n",
    "print(f'image_channels is {image_channels}')\n",
    "print(labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "b17bcc6c-2a1a-4924-988c-5d22b03f7c5d",
   "metadata": {},
   "outputs": [],
   "source": [
    "#定义并训练卷积神经网络"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "3c001a48-79c9-4245-b940-cca3be67edca",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CNN(\n",
      "  (conv1): Conv2d(1, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "  (pool1): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
      "  (conv2): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "  (pool2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
      "  (fc1): Linear(in_features=3136, out_features=100, bias=True)\n",
      "  (fc2): Linear(in_features=100, out_features=10, bias=True)\n",
      "  (softmax): Softmax(dim=1)\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "\n",
    "class CNN(nn.Module):\n",
    "    def __init__(self, image_channels, num_classes):\n",
    "        super(CNN, self).__init__()\n",
    "        \n",
    "        # First convolutional layer\n",
    "        self.conv1 = nn.Conv2d(in_channels=image_channels, out_channels=32, kernel_size=3, padding=1)\n",
    "        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)\n",
    "        \n",
    "        # Second convolutional layer\n",
    "        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)\n",
    "        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)\n",
    "        \n",
    "        # Fully connected layers\n",
    "        self.fc1 = nn.Linear(64 * 7 * 7, 100)  # After two 2x2 max pools, 28x28 -> 7x7\n",
    "        self.fc2 = nn.Linear(100, num_classes)  # 10 classes output\n",
    "\n",
    "        # Softmax\n",
    "        self.softmax = nn.Softmax(dim=1)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        # Remove the reshape operation and directly use x\n",
    "        x = F.relu(self.conv1(x))\n",
    "        x = self.pool1(x)\n",
    "        \n",
    "        # Second conv layer\n",
    "        x = F.relu(self.conv2(x))\n",
    "        x = self.pool2(x)\n",
    "        \n",
    "        # Flatten the output for the fully connected layers\n",
    "        x = x.view(-1, 64 * 7 * 7)\n",
    "        \n",
    "        # Fully connected layers\n",
    "        x = F.relu(self.fc1(x))\n",
    "        x = self.fc2(x)\n",
    "        \n",
    "        # Softmax\n",
    "        x = self.softmax(x)\n",
    "        \n",
    "        return x\n",
    "\n",
    "# Initialize the model\n",
    "model = CNN(image_channels, 10)\n",
    "if torch.cuda.is_available():\n",
    "    model = model.cuda()\n",
    "\n",
    "# Display model architecture\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "1efcab3d-a69d-4e64-aa8e-02778935c721",
   "metadata": {},
   "outputs": [],
   "source": [
    "#使用Adam作为Optimizor训练模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "b0ff7c10-f107-4917-8a3e-b8fbec69a433",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [1/50], Train Loss: 1.5830, CV Loss: 1.5003\n",
      "Epoch [2/50], Train Loss: 1.4885, CV Loss: 1.4919\n",
      "Epoch [3/50], Train Loss: 1.4816, CV Loss: 1.4888\n",
      "Epoch [4/50], Train Loss: 1.4778, CV Loss: 1.4869\n",
      "Epoch [5/50], Train Loss: 1.4755, CV Loss: 1.4881\n",
      "Epoch [6/50], Train Loss: 1.4733, CV Loss: 1.4797\n",
      "Epoch [7/50], Train Loss: 1.4725, CV Loss: 1.4805\n",
      "Epoch [8/50], Train Loss: 1.4714, CV Loss: 1.4817\n",
      "Epoch [9/50], Train Loss: 1.4709, CV Loss: 1.4814\n",
      "Epoch [10/50], Train Loss: 1.4712, CV Loss: 1.4785\n",
      "Epoch [11/50], Train Loss: 1.4704, CV Loss: 1.4767\n",
      "Epoch [12/50], Train Loss: 1.4691, CV Loss: 1.4774\n",
      "Epoch [13/50], Train Loss: 1.4682, CV Loss: 1.4777\n",
      "Epoch [14/50], Train Loss: 1.4685, CV Loss: 1.4773\n",
      "Epoch [15/50], Train Loss: 1.4680, CV Loss: 1.4765\n",
      "Epoch [16/50], Train Loss: 1.4677, CV Loss: 1.4751\n",
      "Epoch [17/50], Train Loss: 1.4673, CV Loss: 1.4765\n",
      "Epoch [18/50], Train Loss: 1.4684, CV Loss: 1.4775\n",
      "Epoch [19/50], Train Loss: 1.4674, CV Loss: 1.4766\n",
      "Epoch [20/50], Train Loss: 1.4671, CV Loss: 1.4769\n",
      "Epoch [21/50], Train Loss: 1.4667, CV Loss: 1.4765\n",
      "Epoch [22/50], Train Loss: 1.4676, CV Loss: 1.4768\n",
      "Epoch [23/50], Train Loss: 1.4665, CV Loss: 1.4755\n",
      "Epoch [24/50], Train Loss: 1.4667, CV Loss: 1.4764\n",
      "Epoch [25/50], Train Loss: 1.4665, CV Loss: 1.4772\n",
      "Epoch [26/50], Train Loss: 1.4669, CV Loss: 1.4804\n",
      "Epoch [27/50], Train Loss: 1.4667, CV Loss: 1.4784\n",
      "Epoch [28/50], Train Loss: 1.4657, CV Loss: 1.4779\n",
      "Epoch [29/50], Train Loss: 1.4659, CV Loss: 1.4761\n",
      "Epoch [30/50], Train Loss: 1.4669, CV Loss: 1.4766\n",
      "Epoch [31/50], Train Loss: 1.4657, CV Loss: 1.4780\n",
      "Epoch [32/50], Train Loss: 1.4654, CV Loss: 1.4777\n",
      "Epoch [33/50], Train Loss: 1.4663, CV Loss: 1.4791\n",
      "Epoch [34/50], Train Loss: 1.4655, CV Loss: 1.4763\n",
      "Epoch [35/50], Train Loss: 1.4664, CV Loss: 1.4788\n",
      "Epoch [36/50], Train Loss: 1.4660, CV Loss: 1.4756\n",
      "Epoch [37/50], Train Loss: 1.4660, CV Loss: 1.4777\n",
      "Epoch [38/50], Train Loss: 1.4660, CV Loss: 1.4783\n",
      "Epoch [39/50], Train Loss: 1.4657, CV Loss: 1.4787\n",
      "Epoch [40/50], Train Loss: 1.4647, CV Loss: 1.4764\n",
      "Epoch [41/50], Train Loss: 1.4660, CV Loss: 1.4772\n",
      "Epoch [42/50], Train Loss: 1.4663, CV Loss: 1.4793\n",
      "Epoch [43/50], Train Loss: 1.4658, CV Loss: 1.4741\n",
      "Epoch [44/50], Train Loss: 1.4665, CV Loss: 1.4767\n",
      "Epoch [45/50], Train Loss: 1.4657, CV Loss: 1.4764\n",
      "Epoch [46/50], Train Loss: 1.4652, CV Loss: 1.4763\n",
      "Epoch [47/50], Train Loss: 1.4648, CV Loss: 1.4776\n",
      "Epoch [48/50], Train Loss: 1.4658, CV Loss: 1.4734\n",
      "Epoch [49/50], Train Loss: 1.4657, CV Loss: 1.4773\n",
      "Epoch [50/50], Train Loss: 1.4661, CV Loss: 1.4763\n"
     ]
    }
   ],
   "source": [
    "# Define loss function and optimizer\n",
    "criterion = nn.CrossEntropyLoss() # Loss\n",
    "optimizer = torch.optim.Adam(model.parameters()) # Adam\n",
    "\n",
    "# Lists to store losses\n",
    "train_losses = []\n",
    "cv_losses = []\n",
    "\n",
    "# Number of epochs\n",
    "num_epochs = 50\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "    model.train()\n",
    "    batch_losses = []\n",
    "    \n",
    "    for batch_x, batch_y in trLoader:\n",
    "        # Forward pass\n",
    "        outputs = model(batch_x)\n",
    "        loss = criterion(outputs, batch_y)\n",
    "        \n",
    "        # Backward pass and optimize\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        \n",
    "        batch_losses.append(loss.item())\n",
    "    \n",
    "    # Calculate average training loss for this epoch\n",
    "    avg_train_loss = sum(batch_losses) / len(batch_losses)\n",
    "    train_losses.append(avg_train_loss)\n",
    "    \n",
    "    # Evaluate on cross-validation set\n",
    "    model.eval()\n",
    "    cv_batch_losses = []\n",
    "    with torch.no_grad():\n",
    "        for cv_x, cv_y in cvLoader:\n",
    "            cv_outputs = model(cv_x)\n",
    "            cv_loss = criterion(cv_outputs, cv_y)\n",
    "            cv_batch_losses.append(cv_loss.item())\n",
    "    \n",
    "    avg_cv_loss = sum(cv_batch_losses) / len(cv_batch_losses)\n",
    "    cv_losses.append(avg_cv_loss)\n",
    "    \n",
    "    print(f'Epoch [{epoch+1}/{num_epochs}], Train Loss: {avg_train_loss:.4f}, CV Loss: {avg_cv_loss:.4f}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "68d22067-9a27-4a52-b864-4d849265018a",
   "metadata": {},
   "outputs": [],
   "source": [
    "#计算识别精度，展示学习曲线"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6a1ca53-be92-4216-becd-1f9f40c35706",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy on training set: 99.54%\n",
      "Accuracy on cross-validation set: 98.49%\n"
     ]
    }
   ],
   "source": [
    "# Calculate and print accuracies for training and cross-validation sets\n",
    "model.eval()\n",
    "with torch.no_grad():\n",
    "    # Training set accuracy\n",
    "    tr_correct = 0\n",
    "    tr_total = 0\n",
    "    for images, labels in trLoader:\n",
    "        outputs = model(images)\n",
    "        _, predicted = torch.max(outputs, 1)\n",
    "        _, true_labels = torch.max(labels, 1)\n",
    "        tr_total += labels.size(0)\n",
    "        tr_correct += (predicted == true_labels).sum().item()\n",
    "    \n",
    "    tr_accuracy = 100 * tr_correct / tr_total\n",
    "    \n",
    "    # Test set accuracy\n",
    "    cv_correct = 0\n",
    "    cv_total = 0\n",
    "    for images, labels in cvLoader:\n",
    "        outputs = model(images)\n",
    "        _, predicted = torch.max(outputs, 1)\n",
    "        _, true_labels = torch.max(labels, 1)\n",
    "        cv_total += labels.size(0)\n",
    "        cv_correct += (predicted == true_labels).sum().item()\n",
    "    \n",
    "    cv_accuracy = 100 * cv_correct / cv_total\n",
    "\n",
    "print(f'Accuracy on training set: {tr_accuracy:.2f}%')\n",
    "print(f'Accuracy on cross-validation set: {cv_accuracy:.2f}%')\n",
    "\n",
    "# Plot training and cross-validation losses\n",
    "plt.figure(figsize=(10, 5))\n",
    "plt.plot(range(1, num_epochs+1), train_losses, label='Training Loss')\n",
    "plt.plot(range(1, num_epochs+1), cv_losses, label='Cross-Validation Loss')\n",
    "plt.xlabel('Epoch')\n",
    "plt.ylabel('Loss')\n",
    "plt.title('Training and Cross-Validation Loss')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
